import matplotlib.pyplot as plt
import numpy as np
import re
import datetime
import matplotlib
import os
# matplotlib.rc("font", family='YouYuan')

def parse_log(log_path, key="loss:", interval=1):
    timer = []
    steps = 0
    losses = []
    with open(log_path, encoding='UTF-8') as f:
        for line in f:
            if "Batches" in line:
                micro = int(re.findall(r"\d+", line.split("Batches")[1])[0])
                epoch = int(line.split("/")[1].split(",")[0])
                steps = micro * epoch
            if key in line:
                try:
                    loss = float(re.findall(r"[+-]?\d+\.?\d+e?[+-]?\d*", line.split(key)[1])[0])
                except:
                    print(line, len(losses))
                    loss = np.mean(losses[-10:-1])
                if loss > 4.0:
                     print(len(losses), loss)
                losses.append(loss)
    losses = np.array(losses)[0: len(losses): interval]
    print(losses)
    elapse = []
    for index, _ in enumerate(timer[:-1]):
        elapse.append((timer[index + 1] - timer[index]).seconds)

    avg_cost = np.round(np.mean(elapse), 2) / 10
    total_cost = (avg_cost * steps) / 60 / 60

    print("==> Processing {} / {}, Total time cost predict {:.2f}h,".format(len(timer) * 10, steps, total_cost), "still need {:.2f}h".format((avg_cost * (steps - len(timer) * 10) / 60 / 60)))

    return losses, 0, 0

def moving_average(data, window_size=10):
    mav = []
    for i in range(len(data) - window_size + 1):
        window = sum(sorted(data[i:i + windiw_size][2:-2]))/(window_size - 4)
        mav.append(window)
    return mav

if __name__ == '__main__':
    paths = ["/src/train25.1.0/MindSpeed-MM/logs/training.log"] # 填写具体需要解析的日志的路径，支持多个日志对比同时解析
    names = [] # 自定义曲线的名称，长度path，可为空
    
    if names is None or len(names) != len(paths):
        names = paths

    all_words = ['actor/kl_loss', 'actor/pg_loss'] # 要解析的参数，需在日志中真实存在，比如["actor/pg_loss", "grpo/score/mean"]

    compare = True
    for keyword in all_words:
        print("*" * 80)
        losses = []
        errors = []
        for i, path in enumerate(paths):
            interval = 2 if "gpu_demo" in names else 1
            print(names[i], end="\n")
            if "demo" in names[i]:
                interval = 1

            l, lower, upper = parse_log(path, keyword, interval=interval)
            losses.append(l)
        print("*" * 80)

        save = True # 是否保存图片
        show = False # 是否展示图片
        need_moving_window = False # 是否做窗口平滑处理
        if need_moving_window:
            losses = [moving_average(loss, window_size=10) for loss in losses]
        plt.figure(figsize=(7,3))

        for path, y_axis in zip(names, losses):
            plt.subplot(1,1,1)
            plt.plot(np.array(list(range(len(y_axis)))), y_axis, label=path)
            plt.axhline(y=0, color='grey', linestyle=':')
            # 可以按照此操作添加上下限制
            if "kl_mean" in keyword:
                plt.ylim((-1, 0.2))
            plt.xlim((0, 2))

            plt.xlabel('training step')
            plt.ylabel(keyword)
            plt.legend(loc='best')
        if save:
            os.makedirs('./plot_images', exist_ok=True)
            plt.savefig(f'./plot_images/{keyword.replace("/", "_")}.png')

        if show:
            plt.show()
            plt.close()
